#include <map>
#include <mpi.h>
#include <math.h>
#include "data.h"
#include "psUtil.h"
#include "psSearch.h"
#include "coloring.h"

using namespace std;

PSSearch::PSSearch(Data &dataObj, ArgStruct &args, int size, int rank) {
	int grpCount = dataObj.getGrpCount(dataObj.getGrpCountVector());

	for (int i = 0; i < grpCount; i++) {
		vector<int> emptyVec;
		for (int j = 0; j < grpCount; j++) 
			emptyVec.push_back(INVALID);

		s.push_back(emptyVec);
	}

	this->alpha = args.alpha;
	this->beta1 = args.beta1;
	this->beta2 = args.beta2;
	this->gamma = args.gamma;

	this->maxColor = args.maxColor;

	dataObj.setDimensions(this->nInds, this->nSteps);

	vector<charVec> data = dataObj.getData();

	for (vector<charVec>::iterator i = data.begin();
	     i != data.end(); i++) {
		vector<int> tempVec;
		for (vector<char>::iterator j = i->begin();
		     j != i->end(); j++) {
			tempVec.push_back(PSUtil::getInt(*j));
		}
		this->data.push_back(tempVec);
	}

        // Initializing MPI datatypes
	this->size = size;
	this->myRank = rank;

        // initialize costMatrix
	for (int i = 0; i < nInds; i++) {
		intVec costRow;
		for (int j = 0; j < nSteps; j++) {
			costRow.push_back(0);
		}
		costMatrix.push_back(costRow);
	}

	allColorPairs = getColorPairs();
}

PSSearch::~PSSearch() {

}

void PSSearch::initializeCosts() {
	for (int i = 0; i < nInds; i++)
		for (int j = 0; j < nSteps; j++)
			costMatrix[i][j] = 0;
}

vector<intIntVec> PSSearch::genAllColoring() {
	Coloring *c = new Coloring(nInds, nSteps,
				   maxColor, data);

	return c->getAllColorings();;
}

double PSSearch::partitionSendReceive(vector<intIntVec> allAssignments, int size) {
	const int nWorkers = size - 1;
	
	double startTime = MPI::Wtime();
	const int nAllColorings = allAssignments.size();

	vector<intIntVec>::iterator iter = allAssignments.begin();

	for (int i = 1; i <= nWorkers; i++) {
		int iPartitionSize = (int) ceil ((float) nAllColorings / nWorkers);
		int lastPartitionSize = nAllColorings % iPartitionSize;
		if ((i == nWorkers - 1) && (lastPartitionSize != 0)) 
			iPartitionSize = lastPartitionSize;

 		MPI::COMM_WORLD.Send(&iPartitionSize,
				     1, MPI::INT, i, TAG_ALL_COLOR_SIZE);

		int sendDataCube[iPartitionSize][nInds][nSteps];
		vector<intIntVec> partition;
		int j = 0;

		while (j < iPartitionSize && iter != allAssignments.end()) {
			partition.push_back(*iter++);
			j++;
		}

		int matrix[nInds][nSteps];
		int coloringNumber = 0;
		for (vector<intIntVec>::iterator x = partition.begin();
		     x != partition.end(); x++) {
			int p = 0;

			for (vector<intVec>::iterator y = x->begin();
			     y != x->end(); y++) {

				int q = 0;
				for (vector<int>::iterator z = y->begin();
				     z != y->end(); z++) {
					sendDataCube[coloringNumber][p][q] = *z;
					q++;
				}
				p++;
			}

			coloringNumber ++;
		}

		MPI::COMM_WORLD.Send(sendDataCube, 
				     iPartitionSize * this->nInds * this->nSteps, 
				     MPI::INT, i, TAG_DATA);
	}

	receiveComputedData();
	double endTime = MPI::Wtime();

	return endTime - startTime;
}

void PSSearch::receiveComputedData() {
	vector<int> minCostVector;
        multimap<int, intMatrixType> matrixGroup;
	int minCost = INVALID;

	for (int workerNumber = 1; 
	     workerNumber <= this->size - 1; workerNumber ++) {
		int minCostAttribs[2];
		int grpMatrix[nInds][nSteps];
		MPI::COMM_WORLD.Recv(minCostAttribs, 2, 
				     MPI::INT, workerNumber, TAG_MIN_COST);
		const int cost = minCostAttribs[0];
		const int numMinCost = minCostAttribs[1];
		minCostVector.push_back(cost);

//		cout << "Node " << workerNumber << "'s Min Cost: " 
//		     << cost << endl;

		for (int k = 0; k < numMinCost; k++) {
			MPI::COMM_WORLD.Recv(grpMatrix, nInds * nSteps, MPI::INT, 
					     workerNumber, TAG_GRP_COST_MATRIX);

			intMatrixType grpMat;

			for (int i = 0; i < nInds; i++) {
				intVec grpRow;
				for (int j = 0; j < nSteps; j++) {
					grpRow.push_back(grpMatrix[i][j]);
				}
				grpMat.push_back(grpRow);
			}

			matrixGroup.insert(make_pair (cost, grpMat));
		}
	}

	int globalMinCost = *(min_element(minCostVector.begin(),
				  minCostVector.end()));
	pair<multimap<int, intMatrixType>::iterator, 
		multimap<int, intMatrixType>::iterator> rSet = 
		matrixGroup.equal_range(globalMinCost);

	for (multimap<int, intMatrixType>::iterator itr = rSet.first;
	     itr != rSet.second; ++itr) {
		cout << "Global Min Cost: " << globalMinCost << endl;
		PSUtil::printMatrix(itr->second);
	}
}

void PSSearch::setSigmaT(intMatrixType coloring) {
	for (intMatrixType::iterator t = coloring.begin();
	     t != coloring.end(); t++){
		intVec tGrp;
		for(intVec::iterator i = t->begin(); i != t->end(); i++) {
			const int curColor = *i;
			if (PSUtil::belongsTo(tGrp, curColor) == false)
				tGrp.push_back(curColor);
		}
		this->sigmaT.push_back(tGrp);
	}
}
bool PSSearch::calcAndSend(int rank) {
	MPI::COMM_WORLD.Recv(&this->partitionSize, 1, MPI::INT, 
			     0, TAG_ALL_COLOR_SIZE);

	int cube[this->partitionSize][nInds][nSteps];
	vector<intIntVec> minCostVectors;
	vector<intIntVec> myPartition;
	vector<int> partitionCosts;

 	MPI::COMM_WORLD.Recv(cube, 
 			     this->partitionSize * this->nInds * this->nSteps,
 			     MPI::INT, 0, TAG_DATA);

	multimap<int, intIntVec> costColoring;

	intMatrixType minGrpCostMatrix;

	for (int p = 0; p < this->partitionSize; p++) {
		vector<intVec> matrix;
		for (int i = 0; i < nInds; i++) {
			vector<int> indColors;
			for (int j = 0; j < nSteps; j++) {
				indColors.push_back(cube[p][i][j]);
			}
			matrix.push_back(indColors);
		}
 		const int cost = getCost(matrix);
		partitionCosts.push_back(cost);
		costColoring.insert(make_pair(cost, matrix));
 	}

	const int minCost = *(min_element(partitionCosts.begin(),
					  partitionCosts.end()));

	pair<multimap<int, intMatrixType>::iterator,
		multimap<int, intMatrixType>::iterator> rSet = 
		costColoring.equal_range(minCost);

	intCubeType lowestCostColorings;
	for(multimap<int, intMatrixType>::iterator itr = rSet.first;
	    itr != rSet.second; ++itr) {
		lowestCostColorings.push_back(itr->second);
	}

	int minCostAttribs[2];
	const int numLowCosts = lowestCostColorings.size();
	minCostAttribs[0] = minCost;
	minCostAttribs[1] = numLowCosts;
 	MPI::COMM_WORLD.Send(minCostAttribs, 2, MPI::INT, 0, TAG_MIN_COST);

	for (intCubeType::iterator itr = lowestCostColorings.begin();
	     itr != lowestCostColorings.end(); itr++) {

		int mgcm[nInds][nSteps]; // Min Group Cost Matrix
		intMatrixType::iterator x = itr->begin();

		for (int i = 0; i < nInds, x != itr->end(); i++, x++) {
			intVec::iterator y = x->begin();

			for (int j = 0; j < nSteps, y != x->end(); 
			     j++, y++) {
				mgcm[i][j] = *y;
			}
		}

		MPI::COMM_WORLD.Send(mgcm, this->nInds * this->nSteps, 
				     MPI::INT, 0, TAG_GRP_COST_MATRIX);

	}
}

vector<ColorPairs> PSSearch::getColorPairs() {
	vector<ColorPairs> allColorPairs;
	const int maxNumColors = this->maxColor;

	for (int i = 1; i <= maxNumColors; i++) {
		map<int, int> colorPairSet;
		for (int j = 1; j <= maxNumColors; j++) {
			ColorPairs cp;
			cp.x = i;
			cp.y = j;
			allColorPairs.push_back(cp);
		}
	}

	return allColorPairs;
}


int PSSearch::computeLowerBound(int i, int t, int fgt) {
	int z = 0;

	for (vector<ColorPairs>::iterator itr = allColorPairs.begin();
	     itr != allColorPairs.end(); itr++) {
		const int x = itr->x;
		const int y = itr->y;

		if (t > 0) {
			costMatrix[i][t] = (int) INFINITY;

			z = costMatrix[i][t - 1];
			if(x != y) z += this->alpha;
			if ((x != fgt) && 
			    (PSUtil::belongsTo(sigmaT[t], x)  == true)) {
				z += this->beta1;
			}
//			if ((x != fgt) && (fgt != epsilon)) 
//				z += this->beta2;

			if (costMatrix[i][t] > z) 
				costMatrix[i][t] = z;

		}
		else {
			costMatrix[i][t] = 0;
			if ((x != fgt) && 
			    (PSUtil::belongsTo(sigmaT[t], x) == true)) {
				z += this->beta1;
			}
			if ((x != fgt) && (fgt != epsilon)) 
				z += this->beta2;
		}
        }
}


int PSSearch::getCost(intMatrixType matrix) {
	int costs[nInds];
	for (int i = 0; i < nInds; i++) 
		costs[i] = 0;

	intMatrixType grpColoring = PSUtil::transposeMat(matrix);
	setSigmaT(grpColoring);
	this->initializeCosts();

	int t = 0;
	for (intMatrixType::iterator cItr = grpColoring.begin();
	     cItr != grpColoring.end(); cItr++) {
		int i = 0;
		for (intVec::iterator rItr = cItr->begin();
		     rItr != cItr->end(); rItr++) {
			computeLowerBound(i, t, *rItr);
			i ++;
		}
		t ++;
	}

	int thisColoringCost = 0;

	for (int i = 0; i < nInds; i++) {
		thisColoringCost += costMatrix[i][nSteps - 1];
	}

	return thisColoringCost;
}



